package com.btl.service.ai.biz.adaptor;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.btl.boot.core.dto.response.R;
import com.btl.service.ai.auth.SecurityContextHolder;
import com.btl.service.ai.biz.dto.AnalysisCustomerQuery;
import com.btl.service.ai.biz.entity.KwDocument;
import com.btl.service.ai.biz.service.IKwDocumentService;
import com.btl.service.ai.support.query.SystemAndUserMessageTransformer;
import com.btl.service.ai.support.rag.embedding.strategy.EmbeddingModelTypeEnum;
import com.btl.service.ai.support.rag.embedding.strategy.RetrieverStrategy;
import com.btl.service.ai.support.rag.embedding.strategy.RetrieverStrategyFactory;
import com.btl.service.ai.support.rag.generation.GenerationModelEnum;
import com.btl.service.ai.support.rag.generation.assistant.GenerationAssistantFactory;
import com.btl.service.ai.support.rag.generation.assistant.GenerationAssistantStrategy;
import com.btl.service.ai.support.rag.generation.knowledge.GenerationKnowledgeFactory;
import com.btl.service.ai.support.rag.generation.knowledge.GenerationKnowledgeStrategy;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.Result;
import dev.langchain4j.store.embedding.filter.Filter;
import dev.langchain4j.store.embedding.filter.comparison.IsIn;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.annotation.Resource;
import jakarta.validation.Valid;
import jakarta.validation.constraints.NotBlank;
import jakarta.validation.constraints.NotNull;
import lombok.extern.slf4j.Slf4j;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import org.springframework.stereotype.Controller;

import java.util.ArrayList;
import java.util.List;

/**
 * <p>
 * 前端控制器
 * </p>
 *
 * @author Tim
 * @since 2025-07-22
 */
@Slf4j
@Validated
@RestController
@Tag(name = "助手")
@Controller
@RequestMapping("agent")
public class AiAgentController {

    @Resource
    private IKwDocumentService iKwDocumentService;

    /**
     * 类型提问
     *
     * @param analysisCustomerQuery 查询条件
     * @return List<SplitSegmentRsp>
     */
    @PostMapping(value = "/analysis-customer")
    public R<String> analysis(@RequestBody @Valid AnalysisCustomerQuery analysisCustomerQuery) {
        String messages = analysisCustomerQuery.getMessages();
        Query query = new Query(analysisCustomerQuery.getQuestion());
        SystemAndUserMessageTransformer queryTransformer = new SystemAndUserMessageTransformer(messages);
        SystemAndUserMessageTransformer.Message transform = queryTransformer.transform(query);
        Long memoryId = analysisCustomerQuery.getMemoryId();
        GenerationAssistantStrategy strategy = GenerationAssistantFactory.getStrategy(GenerationModelEnum.OPEN_AI_CHAT_MODEL);
        if (analysisCustomerQuery.getFormatJson()) {
            Result<String> generateJson = strategy.generateJson(memoryId, transform.systemMessage(), transform.userMessage(), query, String.class);
            return R.ok(generateJson.content());
        }
        Result<String> generate = strategy.generate(memoryId, transform.systemMessage(), transform.userMessage(), query);
        return R.ok(generate.content());
    }

    /**
     * 增强生成
     *
     * @param question 查询条件
     * @return List<SplitSegmentRsp>
     */
    @GetMapping(value = "/generate")
    public R<String> generate(@RequestParam(value = "question", required = false) @NotNull @NotBlank(message = "问题不能为空") String question,
                              @RequestParam(value = "memoryId", required = false) @NotNull(message = "记忆ID不能为空") Long memoryId) {
        Long childId = SecurityContextHolder.get()
                .getChildId();
        List<KwDocument> kwDocuments = iKwDocumentService.list(new LambdaQueryWrapper<KwDocument>().eq(KwDocument::getChildId, childId));
        GenerationKnowledgeStrategy generationStrategy = GenerationKnowledgeFactory.getStrategy(GenerationModelEnum.OPEN_AI_CHAT_MODEL);
        if (!kwDocuments.isEmpty()) {
            String collectionName = "kw_segment";
            RetrieverStrategy retrieverStrategy = RetrieverStrategyFactory.retrieverStrategy(EmbeddingModelTypeEnum.LOCAL_BGE_SMALL_ZH_V15_EMBEDDING_MODEL);
            Filter filter = new IsIn("documentId", kwDocuments.stream()
                    .map(KwDocument::getId)
                    .map(Object::toString)
                    .toList());
            List<Content> contents = retrieverStrategy.retriever(collectionName, filter, Query.from(question));
            Result<String> generate = generationStrategy.generate(memoryId, contents, new Query(question));
            return R.ok(generate.content());
        }
        Result<String> generate = generationStrategy.generate(memoryId, new ArrayList<>(), new Query(question));
        return R.ok(generate.content());
    }

}
